import numpy as np
from typing import Union, List
import pandas as pd

class strategy:
    """策略类，包含各种数据处理和约束函数"""
    
    @staticmethod
    def A95_lower_limit(df: Union[pd.DataFrame, np.ndarray], d_std: float) -> float:
        """
        计算除0以外的A95下限
        
        Args:
            df: 输入数据（DataFrame或数组）
            d_std: 标准差倍数（0-1之间的小数）
            
        Returns:
            float: A95下限值
        """
        # 转换为数组并展平
        data = np.array(df).flatten()
        
        # 使用布尔索引过滤非零值，比循环更高效
        non_zero_data = data[data != 0]
        
        # 如果所有值都是0，返回0
        if len(non_zero_data) == 0:
            return 0.0
            
        # 降序排序并计算位置
        sorted_data = np.sort(non_zero_data)[::-1]  # 降序排序
        index = int(len(sorted_data) * d_std)
        
        # 确保索引在有效范围内
        index = max(0, min(index, len(sorted_data) - 1))
        
        return float(sorted_data[index])
    
    @staticmethod
    def Nonnegative_number(y_pred: Union[List, np.ndarray]) -> np.ndarray:
        """
        将负数转换为0，确保没有负数
        
        Args:
            y_pred: 预测值数组
            
        Returns:
            np.ndarray: 处理后的非负数组
        """
        # 转换为numpy数组并使用向量化操作
        y_array = np.array(y_pred)
        
        # 使用numpy的where函数，比列表推导式更高效
        return np.where(y_array >= 0, y_array, 0)
    
    @staticmethod
    def boll(y_pred: Union[List, np.ndarray], multiple: float) -> np.ndarray:
        """
        Boll标准差约束，将超出均值±倍数标准差的值进行截断
        
        Args:
            y_pred: 预测值数组
            multiple: 标准差倍数
            
        Returns:
            np.ndarray: 约束后的数组
        """
        # 转换为numpy数组
        y_array = np.array(y_pred)
        
        # 计算均值和标准差
        mean_val = np.mean(y_array)
        std_val = np.std(y_array)
        
        # 计算上下界（修复原代码中的错误：z12应该是减去multiple*std_val）
        upper_bound = mean_val + multiple * std_val
        lower_bound = mean_val - multiple * std_val
        
        # 使用numpy的clip函数进行截断，比循环更高效
        return np.clip(y_array, lower_bound, upper_bound)